import pyautogui
import imageio
import os
import time
import numpy as np
import tkinter as tk
from tkinter import filedialog
import threading

# 设置默认参数
fps = 25  # 帧率
output_directory = os.getcwd()  # 默认输出目录
output_filename = 'screen_recording.mp4'  # 文件名
output_path = os.path.join(output_directory, output_filename)  # 完整路径

recording = False  # 是否正在录制的标志
video_writer = None  # 视频写入对象
record_thread = None  # 录制线程

def start_recording():
    global recording, video_writer, output_path, record_thread
    if not recording:
        output_directory = entry.get()
        if not os.path.exists(output_directory):
            os.makedirs(output_directory)
        output_path = os.path.join(output_directory, output_filename)
        video_writer = imageio.get_writer(output_path, fps=fps)
        recording = True
        enable_buttons(False, True)  # 开始录制后，开始按钮无效，停止按钮有效
        record_thread = threading.Thread(target=record_screen)
        record_thread.start()

def stop_recording():
    global recording, video_writer, record_thread
    if recording:
        recording = False
        if video_writer is not None:
            video_writer.close()
            video_writer = None
        if record_thread is not None and record_thread.is_alive():
            record_thread.join()
        enable_buttons(True, False)  # 停止录制后，开始按钮有效，停止按钮无效

def record_screen():
    while recording:
        screenshot = pyautogui.screenshot()
        frame = np.array(screenshot)
        video_writer.append_data(frame)
        time.sleep(1 / fps)

def select_output_directory():
    directory = filedialog.askdirectory(initialdir=output_directory)
    if directory:
        entry.delete(0, tk.END)
        entry.insert(0, directory)

def enable_buttons(start_enabled, stop_enabled):
    start_button.config(state=tk.NORMAL if start_enabled else tk.DISABLED)
    stop_button.config(state=tk.NORMAL if stop_enabled else tk.DISABLED)

# 创建主窗口
root = tk.Tk()
root.title("屏幕录制工具")
root.geometry("250x150")  # 设置窗口大小

# 路径选择框
entry = tk.Entry(root, width=40)
entry.insert(0, output_directory)
entry.pack(pady=5)

# 选择路径按钮
select_button = tk.Button(root, text="选择输出目录", command=select_output_directory, width=20)
select_button.pack(pady=2)

# 开始录制按钮
start_button = tk.Button(root, text="开始录制", command=start_recording, width=20)
start_button.pack(pady=2)

# 停止录制按钮
stop_button = tk.Button(root, text="停止录制", command=stop_recording, width=20)
stop_button.pack(pady=2)

# 初始状态下，停止按钮无效
enable_buttons(True, False)

# 运行主循环
root.mainloop()

